import tushare as ts
import pymysql
from dataclasses import dataclass
from typing import List, Optional
import time
from datetime import datetime

# 股票基本信息数据类
@dataclass
class StockBasic:
    """股票基本信息"""
    ts_code: str        # TS股票代码
    symbol: str         # 股票代码
    name: str           # 股票名称
    area: str           # 所在地域
    industry: str       # 所属行业
    list_date: str      # 上市日期

class MysqlDatabase:
    """MySQL数据库操作类"""
    
    def __init__(self, host: str, port: int, user: str, 
                 password: str, database: str):
        """初始化数据库连接"""
        self.conn = pymysql.connect(
            host=host,
            port=port,
            user=user,
            password=password,
            database=database,
            charset='utf8mb4',
            cursorclass=pymysql.cursors.DictCursor
        )
        self._init_tables()
    
    def _init_tables(self):
        """初始化股票基本信息表"""
        sql = """
        CREATE TABLE IF NOT EXISTS stock_basic (
            ts_code VARCHAR(20) PRIMARY KEY,
            symbol VARCHAR(10) NOT NULL,
            name VARCHAR(50) NOT NULL,
            area VARCHAR(50),
            industry VARCHAR(50),
            list_date VARCHAR(255)
        ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4
        """
        self._execute_sql(sql)
    
    def _execute_sql(self, sql: str, params=None):
        """执行SQL语句"""
        try:
            print(f"执行SQL: {sql}")
            print(f"参数: {params}")
            
            with self.conn.cursor() as cursor:
                cursor.executemany(sql, params)
            self.conn.commit()
            return True
        except Exception as e:
            print(f"执行SQL失败: {e}")
            print(f"失败SQL: {sql}")
            print(f"失败参数: {params}")
            self.conn.rollback()
            return False
    
    def save_stock_basic(self, stocks: List[StockBasic]) -> bool:
        """保存股票基本信息到数据库"""
        if not stocks:
            return False
            
        sql = """
        INSERT INTO stock_basic 
        (ts_code, symbol, name, area, industry, list_date)
        VALUES (%s, %s, %s, %s, %s, %s)
        """
        
        # now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        data = [
            (stock.ts_code, stock.symbol, stock.name,
             stock.area, stock.industry, stock.list_date)
            for stock in stocks
        ]
        data1 = data
        return self._execute_sql(sql, data1)

def get_tushare_stocks(token: str) -> List[StockBasic]:
    """从Tushare获取股票基本信息"""
    print("正在从Tushare获取股票基本信息...")
    
    # 设置Tushare token
    ts.set_token(token)
    pro = ts.pro_api()
    
    try:
        # 获取上市股票数据
        df = pro.stock_basic(
            exchange='', 
            list_status='L', 
            fields='ts_code,symbol,name,area,industry,list_date'
        )
        
        # 将DataFrame转换为StockBasic对象列表
        # 1. 使用iterrows()遍历DataFrame的每一行
        # 2. 将每行数据转换为StockBasic对象
        # 3. 最终得到一个包含所有股票信息的列表
        stocks = []
        for _, row in df.iterrows():
            stock = StockBasic(
                ts_code=row['ts_code'],    # 股票代码(带交易所后缀)
                symbol=row['symbol'],      # 股票代码(无后缀)
                name=row['name'],          # 股票名称
                area=row['area'],          # 地区
                industry=row['industry'],  # 行业
                list_date=row['list_date'] # 上市日期
            )
            stocks.append(stock)
        
        print(f"成功获取 {len(stocks)} 条股票数据")
        return stocks
        
    except Exception as e:
        print(f"获取Tushare数据失败: {e}")
        return []

def main():
    """主函数"""
    # 配置信息 - 请根据实际情况修改
    config = {
        'tushare_token': 'f9070562c743804c9627e56832d30c32a22f2f8c06633357428e6d98',  # 替换为你的Tushare token
        'db_host': 'localhost',
        'db_port': 3306,
        'db_user': 'root',
        'db_password': '12345678',
        'db_name': 'vnpy'
    }
    
    # 获取数据
    stocks = get_tushare_stocks(config['tushare_token'])
    
    if stocks:
        # 初始化数据库
        db = MysqlDatabase(
            host=config['db_host'],
            port=config['db_port'],
            user=config['db_user'],
            password=config['db_password'],
            database=config['db_name']
        )
        
        # 保存数据
        if db.save_stock_basic(stocks):
            print("股票数据保存成功")
        else:
            print("股票数据保存失败")
    else:
        print("没有获取到股票数据")

if __name__ == "__main__":
    main()
